import tensorflow as tf
import numpy as np
from new import get_var_wrap
#import tensorflow.linalg
from kron import kron


def ht1(inp,
        inp_modes,#n
		out1_modes,
        out_modes,
		matin_ranks,  #rk
		matout_ranks,  # rk
		cores_initializer=tf.contrib.layers.xavier_initializer(uniform=False),
#		cores_initializer=tf.truncated_normal_initializer(stddev=0.1),
        cores_regularizer=None,
#       biases_initializer=tf.constant_initializer(0.1),
        biases_regularizer=None,
        trainable=True,
        cpu_variables=False,
		scope='a'):


#	with tf.variable_scope(scope,reuse=tf.AUTO_REUSE):
	with tf.variable_scope(scope):
		out = inp
		dimin = len(inp_modes)
		dimout = len(out1_modes)
		matin_cores = []
		matout_cores = []
		transtensorin = []
		transtensorout = []
		trans_kronin = []
		trans_kronout = []
		matin_kron = []
		matout_kron = []
		cinit = cores_initializer
		creg = cores_regularizer
		for i in range(dimin):
			cinit = cores_initializer
			creg = cores_regularizer
			matin_cores.append(get_var_wrap('mat_corein_%d' % (i + 1),
                                            shape=[inp_modes[i], matin_ranks[i]],
                                            initializer=cinit,
                                            regularizer=creg,
                                            trainable=trainable,
                                            cpu_variable=cpu_variables))
		
		for i in range(0,dimin-1,2):
			matin_kron.append(kron(matin_cores[i],matin_cores[i+1]))#Ua与Ub的克罗内克积
		for i in range(dimout):
			cinit = cores_initializer
			creg = cores_regularizer
			matout_cores.append(get_var_wrap('mat_coreout_%d' % (i + 1),
											shape=[out1_modes[i], matout_ranks[i]],
											initializer=cinit,
											regularizer=creg,
											trainable=trainable,
											cpu_variable=cpu_variables))
		
		for i in range(0,dimout - 1,2):
			matout_kron.append(kron(matout_cores[i], matout_cores[i + 1]))  # Ua与Ub的克罗内克积

		blast1 = get_var_wrap('mat_coreinlast',
                                            shape=[matin_ranks[-1], 1],
                                            initializer=cinit,
                                            regularizer=creg,
                                            trainable=trainable,
                                            cpu_variable=cpu_variables)
		blast2 = get_var_wrap('mat_coreoutlast',
							  shape=[matout_ranks[-1], 1],
							  initializer=cinit,
							  regularizer=creg,
							  trainable=trainable,
							  cpu_variable=cpu_variables)
		n=0
		sdim=dimin
		while sdim/2>=1:
			sdim=sdim/2
			n=n+1
		for i in range(n):
			cinit = cores_initializer
			creg = cores_regularizer
			if i==0:
				n1=0
			else:
				n1=int(2*dimin*(1-1/(2**i)))
			n2=int(2*dimin*(1-1/(2**(i+1))))

			for j in range(2**(n-1-i)):
				transtensorin.append(get_var_wrap('transtensorin_%d' % (10*i+j),
                                                shape=[matin_ranks[n1+2*j]*matin_ranks[n1+2*j+1],matin_ranks[n2+j]],
                                                initializer=cinit,
                                                regularizer=creg,
                                                trainable=trainable,
                                                cpu_variable=cpu_variables))#Uab,Uabcd

		for i in range(dimin//2):
			trans_kronin.append(tf.matmul(matin_kron[i],transtensorin[i]))
		for i in range((dimin//2)//2):
			fullin = kron(trans_kronin[i],trans_kronin[i+1])
		fullin = tf.matmul(fullin,transtensorin[-1])


		for i in range(n):
			cinit = cores_initializer
			creg = cores_regularizer
			if i == 0:
				n1 = 0
			else:
				n1 = int(2 * dimin * (1 - 1 / (2 ** i)))
			n2 = int(2 * dimin * (1 - 1 / (2 ** (i + 1))))

			for j in range(2 ** (n - 1 - i)):
				transtensorout.append(get_var_wrap('transtensorout_%d' % (10 * i + j),
												  shape=[matout_ranks[n1 + 2 * j] * matout_ranks[n1 + 2*j + 1],
														 matout_ranks[n2 + j]],
												  initializer=cinit,
												  regularizer=creg,
												  trainable=trainable,
												  cpu_variable=cpu_variables))  # Uab,Uabcd

		for i in range(dimin//2):
			trans_kronout.append(tf.matmul(matout_kron[i],transtensorout[i]))
		for i in range((dimin//2)//2):
			fullout = kron(trans_kronout[i],trans_kronout[i+1])
		fullout = tf.matmul(fullout,transtensorout[-1])

		fullin = tf.matmul(fullin,blast1)
		fullout = tf.matmul(fullout,blast2)

		out = tf.matmul(out,fullin)
		fullout = tf.transpose(fullout,[1,0])
		out = tf.matmul(out,fullout)


		#if biases_initializer is not None:

		#	biases = get_var_wrap('biases',
  #                                shape=out_modes,
  #                                initializer=biases_initializer,
  #                                regularizer=biases_regularizer,
  #                                trainable=trainable,
  #                                cpu_variable=cpu_variables)

  #          #out = tf.add(tf.reshape(out, [ np.prod(out_modes),-1]), biases, name="out")
		#	out = tf.add(out, biases, name="out")
		#else:
		#	out = tf.reshape(out, [ np.prod(out_modes),-1], name="out")
		biases = tf.Variable(tf.constant(0.1, shape=out_modes))
		out = out + biases
	return [blast1,out]
